from datetime import datetime, timezone
from unittest.mock import patch
from uuid import UUID

import pytest
from freezegun import freeze_time
from httpx import AsyncClient
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from dstack._internal.core.models.projects import ProjectRole
from dstack._internal.core.models.runs import JobStatus, RunStatus
from dstack._internal.core.models.users import GlobalRole
from dstack._internal.server.models import MemberModel, UserModel
from dstack._internal.server.services.projects import add_project_member
from dstack._internal.server.testing.common import (
    create_job,
    create_probe,
    create_project,
    create_repo,
    create_run,
    create_user,
    get_auth_headers,
)


class TestListUsers:
    @pytest.mark.asyncio
    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
    async def test_returns_40x_if_not_authenticated(self, test_db, client: AsyncClient):
        response = await client.post("/api/users/list")
        assert response.status_code in [401, 403]

    @pytest.mark.asyncio
    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
    async def test_admins_see_all_non_deleted_users(
        self, test_db, session: AsyncSession, client: AsyncClient
    ):
        admin = await create_user(
            session=session,
            name="admin",
            created_at=datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc),
            global_role=GlobalRole.ADMIN,
        )
        other_user = await create_user(
            session=session,
            name="other_user",
            created_at=datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc),
            global_role=GlobalRole.USER,
        )
        await create_user(
            session=session,
            name="deleted_user",
            created_at=datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc),
            global_role=GlobalRole.USER,
            deleted=True,
        )
        response = await client.post("/api/users/list", headers=get_auth_headers(admin.token))
        assert response.status_code in [200]
        assert response.json() == [
            {
                "id": str(admin.id),
                "username": admin.name,
                "created_at": "2023-01-02T03:04:00+00:00",
                "global_role": admin.global_role,
                "email": None,
                "active": True,
                "permissions": {
                    "can_create_projects": True,
                },
                "ssh_public_key": None,
            },
            {
                "id": str(other_user.id),
                "username": other_user.name,
                "created_at": "2023-01-02T03:04:00+00:00",
                "global_role": other_user.global_role,
                "email": None,
                "active": True,
                "permissions": {
                    "can_create_projects": True,
                },
                "ssh_public_key": None,
            },
        ]

    @pytest.mark.asyncio
    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
    async def test_non_admins_see_only_themselves(
        self, test_db, session: AsyncSession, client: AsyncClient
    ):
        await create_user(
            session=session,
            name="admin",
            created_at=datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc),
            global_role=GlobalRole.ADMIN,
        )
        other_user = await create_user(
            session=session,
            name="other_user",
            created_at=datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc),
            global_role=GlobalRole.USER,
        )
        response = await client.post("/api/users/list", headers=get_auth_headers(other_user.token))
        assert response.status_code in [200]
        assert response.json() == [
            {
                "id": str(other_user.id),
                "username": other_user.name,
                "created_at": "2023-01-02T03:04:00+00:00",
                "global_role": other_user.global_role,
                "email": None,
                "active": True,
                "permissions": {
                    "can_create_projects": True,
                },
                "ssh_public_key": None,
            }
        ]


class TestGetMyUser:
    @pytest.mark.asyncio
    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
    async def test_returns_40x_if_not_authenticated(self, test_db, client: AsyncClient):
        response = await client.post("/api/users/get_my_user")
        assert response.status_code in [401, 403]

    @pytest.mark.asyncio
    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
    async def test_returns_40x_if_deactivated(
        self, test_db, session: AsyncSession, client: AsyncClient
    ):
        user = await create_user(session=session, active=False)
        response = await client.post(
            "/api/users/get_my_user", headers=get_auth_headers(user.token)
        )
        assert response.status_code in [401, 403]
        user.active = True
        await session.commit()
        response = await client.post(
            "/api/users/get_my_user", headers=get_auth_headers(user.token)
        )
        assert response.status_code == 200

    @pytest.mark.asyncio
    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
    async def test_returns_logged_in_user(
        self, test_db, session: AsyncSession, client: AsyncClient
    ):
        user = await create_user(
            session=session,
            created_at=datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc),
            ssh_public_key="public-key",
            ssh_private_key="private-key",
        )
        response = await client.post(
            "/api/users/get_my_user", headers=get_auth_headers(user.token)
        )
        assert response.status_code == 200
        assert response.json() == {
            "id": str(user.id),
            "username": user.name,
            "created_at": "2023-01-02T03:04:00+00:00",
            "global_role": user.global_role,
            "email": None,
            "creds": {"token": user.token.get_plaintext_or_error()},
            "active": True,
            "permissions": {
                "can_create_projects": True,
            },
            "ssh_private_key": "private-key",
            "ssh_public_key": "public-key",
        }

    @pytest.mark.asyncio
    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
    async def test_generates_ssh_key_if_missing(
        self, test_db, session: AsyncSession, client: AsyncClient
    ):
        user = await create_user(
            session=session,
            ssh_public_key=None,
            ssh_private_key=None,
        )
        with patch("dstack._internal.utils.crypto.generate_rsa_key_pair_bytes") as gen_mock:
            gen_mock.return_value = (b"private-key", b"ssh-rsa AAA.public-key user\n")
            response = await client.post(
                "/api/users/get_my_user", headers=get_auth_headers(user.token)
            )
        assert response.status_code == 200
        data = response.json()
        assert data["ssh_private_key"] == "private-key"
        assert data["ssh_public_key"] == "ssh-rsa AAA.public-key user\n"
        await session.refresh(user)
        assert user.ssh_private_key == data["ssh_private_key"]
        assert user.ssh_public_key == data["ssh_public_key"]


class TestGetUser:
    @pytest.mark.asyncio
    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
    async def test_returns_40x_if_not_authenticated(self, test_db, client: AsyncClient):
        response = await client.post("/api/users/get_user")
        assert response.status_code in [401, 403]

    @pytest.mark.asyncio
    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
    async def test_returns_400_if_not_global_admin(
        self, test_db, session: AsyncSession, client: AsyncClient
    ):
        user = await create_user(session=session, global_role=GlobalRole.USER)
        other_user = await create_user(session=session, name="other_user", token="1234")
        response = await client.post(
            "/api/users/get_user",
            headers=get_auth_headers(user.token),
            json={"username": other_user.name},
        )
        assert response.status_code == 400

    @pytest.mark.asyncio
    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
    async def test_returns_logged_in_user(
        self, test_db, session: AsyncSession, client: AsyncClient
    ):
        user = await create_user(session=session, global_role=GlobalRole.ADMIN)
        other_user = await create_user(
            session=session,
            name="other_user",
            token="1234",
            created_at=datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc),
        )
        response = await client.post(
            "/api/users/get_user",
            headers=get_auth_headers(user.token),
            json={"username": other_user.name},
        )
        assert response.status_code == 200
        assert response.json() == {
            "id": str(other_user.id),
            "username": other_user.name,
            "created_at": "2023-01-02T03:04:00+00:00",
            "global_role": other_user.global_role,
            "email": None,
            "creds": {"token": "1234"},
            "active": True,
            "permissions": {
                "can_create_projects": True,
            },
            "ssh_private_key": None,
            "ssh_public_key": None,
        }


class TestCreateUser:
    @pytest.mark.asyncio
    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
    async def test_returns_40x_if_not_authenticated(self, test_db, client: AsyncClient):
        response = await client.post("/api/users/create")
        assert response.status_code in [401, 403]

    @pytest.mark.asyncio
    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
    @freeze_time(datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc))
    async def test_creates_user(self, test_db, session: AsyncSession, client: AsyncClient):
        user = await create_user(name="admin", session=session)
        with patch("uuid.uuid4") as uuid_mock:
            uuid_mock.return_value = UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e")
            response = await client.post(
                "/api/users/create",
                headers=get_auth_headers(user.token),
                json={
                    "username": "test",
                    "global_role": GlobalRole.USER,
                    "email": "test@example.com",
                    "active": True,
                },
            )
        assert response.status_code == 200
        user_data = response.json()
        ssh_public_key = user_data["ssh_public_key"]
        assert user_data == {
            "id": "1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e",
            "username": "test",
            "created_at": "2023-01-02T03:04:00+00:00",
            "global_role": "user",
            "email": "test@example.com",
            "active": True,
            "permissions": {
                "can_create_projects": True,
            },
            "ssh_public_key": ssh_public_key,
        }
        res = await session.execute(select(UserModel).where(UserModel.name == "test"))
        assert len(res.scalars().all()) == 1

    @pytest.mark.asyncio
    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
    @freeze_time(datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc))
    async def test_return_400_if_username_taken(
        self, test_db, session: AsyncSession, client: AsyncClient
    ):
        user = await create_user(
            name="admin",
            session=session,
        )
        with patch("uuid.uuid4") as uuid_mock:
            uuid_mock.return_value = UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e")
            response = await client.post(
                "/api/users/create",
                headers=get_auth_headers(user.token),
                json={
                    "username": "Test",
                    "global_role": GlobalRole.USER,
                },
            )
        assert response.status_code == 200
        user_data = response.json()
        ssh_public_key = user_data["ssh_public_key"]
        assert user_data == {
            "id": "1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e",
            "username": "Test",
            "created_at": "2023-01-02T03:04:00+00:00",
            "global_role": "user",
            "email": None,
            "active": True,
            "permissions": {
                "can_create_projects": True,
            },
            "ssh_public_key": ssh_public_key,
        }
        # Username uniqueness check should be case insensitive
        for username in ["test", "Test", "TesT"]:
            with patch("uuid.uuid4") as uuid_mock:
                uuid_mock.return_value = UUID("1b0e1b45-2f8c-4ab6-8010-a0d1a3e44e0e")
                response = await client.post(
                    "/api/users/create",
                    headers=get_auth_headers(user.token),
                    json={
                        "username": username,
                        "global_role": GlobalRole.USER,
                    },
                )
            assert response.status_code == 400
        res = await session.execute(
            select(UserModel).where(UserModel.name.in_(["test", "Test", "TesT"]))
        )
        assert len(res.scalars().all()) == 1

    @pytest.mark.asyncio
    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
    @freeze_time(datetime(2023, 1, 2, 3, 4, tzinfo=timezone.utc))
    async def test_returns_400_if_username_invalid(
        self,
        test_db,
        session: AsyncSession,
        client: AsyncClient,
    ):
        user = await create_user(
            name="admin",
            session=session,
        )
        response = await client.post(
            "/api/users/create",
            headers=get_auth_headers(user.token),
            json={
                "username": "Invalid#$username",
                "global_role": GlobalRole.USER,
            },
        )
        assert response.status_code == 400


class TestDeleteUsers:
    @pytest.mark.asyncio
    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
    async def test_returns_40x_if_not_authenticated(self, test_db, client: AsyncClient):
        response = await client.post("/api/users/delete")
        assert response.status_code in [401, 403]

    @pytest.mark.asyncio
    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
    @pytest.mark.parametrize("username", ["test", "a" * 50])
    async def test_deletes_users(
        self, test_db, session: AsyncSession, client: AsyncClient, username: str
    ):
        admin = await create_user(name="admin", session=session)
        user = await create_user(name=username, session=session)
        response = await client.post(
            "/api/users/delete",
            headers=get_auth_headers(admin.token),
            json={"users": [user.name]},
        )
        assert response.status_code == 200
        res = await session.execute(select(UserModel).where(UserModel.name == user.name))
        assert len(res.scalars().all()) == 0

    @pytest.mark.asyncio
    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
    async def test_returns_400_if_users_not_exist(
        self, test_db, session: AsyncSession, client: AsyncClient
    ):
        admin = await create_user(name="admin", session=session)
        user1 = await create_user(name="test1", session=session)
        user2 = await create_user(name="test2", session=session)
        response = await client.post(
            "/api/users/delete",
            headers=get_auth_headers(admin.token),
            json={"users": [user1.name, "non_existing_user"]},
        )
        assert response.status_code == 400
        response = await client.post(
            "/api/users/delete",
            headers=get_auth_headers(admin.token),
            json={"users": [user1.name, user2.name]},
        )
        assert response.status_code == 200

    @pytest.mark.asyncio
    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
    @pytest.mark.usefixtures("image_config_mock")
    async def test_deletes_user_with_resources(
        self, test_db, session: AsyncSession, client: AsyncClient
    ):
        admin = await create_user(name="admin", session=session)
        user = await create_user(name="temp", session=session, global_role=GlobalRole.USER)
        project = await create_project(session=session, owner=user)
        repo = await create_repo(session=session, project_id=project.id)
        run = await create_run(
            session=session,
            project=project,
            repo=repo,
            user=user,
            status=RunStatus.RUNNING,
        )
        job = await create_job(session=session, run=run, status=JobStatus.RUNNING)
        await create_probe(session=session, job=job)
        response = await client.post(
            "/api/users/delete",
            headers=get_auth_headers(admin.token),
            json={"users": [user.name]},
        )
        assert response.status_code == 200
        res = await session.execute(select(UserModel).where(UserModel.name == user.name))
        assert res.scalar() is None

    @pytest.mark.asyncio
    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
    @pytest.mark.usefixtures("image_config_mock")
    async def test_deleting_users_deletes_members(
        self, test_db, session: AsyncSession, client: AsyncClient
    ):
        admin = await create_user(name="admin", session=session)
        user = await create_user(name="temp", session=session, global_role=GlobalRole.USER)
        project = await create_project(session=session, owner=user)
        await add_project_member(
            session=session, project=project, user=user, project_role=ProjectRole.USER
        )
        response = await client.post(
            "/api/users/delete",
            headers=get_auth_headers(admin.token),
            json={"users": [user.name]},
        )
        assert response.status_code == 200
        res = await session.execute(select(UserModel).where(UserModel.name == user.name))
        assert res.scalar() is None
        res = await session.execute(select(MemberModel).where(MemberModel.user_id == user.id))
        assert res.scalar() is None


class TestRefreshToken:
    @pytest.mark.asyncio
    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
    async def test_returns_40x_if_not_authenticated(self, test_db, client: AsyncClient):
        response = await client.post("/api/users/refresh_token")
        assert response.status_code in [401, 403]

    @pytest.mark.asyncio
    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
    async def test_refreshes_token(self, test_db, session: AsyncSession, client: AsyncClient):
        user1 = await create_user(name="user1", session=session)
        old_token = user1.token
        response = await client.post(
            "/api/users/refresh_token",
            headers=get_auth_headers(user1.token),
            json={"username": user1.name},
        )
        assert response.status_code == 200
        assert response.json()["creds"]["token"] != old_token
        await session.refresh(user1)
        assert user1.token != old_token

    @pytest.mark.asyncio
    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
    async def test_returns_403_if_non_admin_refreshes_for_other_user(
        self, test_db, session: AsyncSession, client: AsyncClient
    ):
        user1 = await create_user(name="user1", session=session, global_role=GlobalRole.USER)
        user2 = await create_user(name="user2", session=session)
        response = await client.post(
            "/api/users/refresh_token",
            headers=get_auth_headers(user1.token),
            json={"username": user2.name},
        )
        assert response.status_code == 403

    @pytest.mark.asyncio
    @pytest.mark.parametrize("test_db", ["sqlite", "postgres"], indirect=True)
    async def test_global_admin_refreshes_token(
        self, test_db, session: AsyncSession, client: AsyncClient
    ):
        user1 = await create_user(name="user1", session=session, global_role=GlobalRole.ADMIN)
        user2 = await create_user(name="user2", session=session)
        old_token = user2.token
        response = await client.post(
            "/api/users/refresh_token",
            headers=get_auth_headers(user1.token),
            json={"username": user2.name},
        )
        assert response.status_code == 200
        assert response.json()["creds"]["token"] != old_token
        await session.refresh(user2)
        assert user2.token != old_token
